import weakref
from bisect import bisect
from typing import Callable, Iterable, Iterator, Optional, Union
from warnings import warn

import matplotlib.pyplot as plt
import numpy as np
from IPython.display import display
from ipywidgets import widgets

from .live_channel import InputType, IntervalLiveChannel, IntervalLiveProducer, OutputType
from .live_figure import LiveFigure


class LinesBufferChannel(IntervalLiveChannel):
    def __init__(self, *, interval: float = 1 / 30, win: float = 200, bins: int = 10):
        super().__init__(interval=interval)
        self._win = win
        self._bins = bins

        # buffers
        self._xys_logs_buffer = bytearray()
        self._ys_n = None
        self._cursor = 0

    def clear(self):
        self.send(None)

    def _update_buffer(self, update: InputType):
        if update is None:  # clear
            self._xys_logs_buffer.clear()
            self._ys_n = None
            self._cursor = 0
            return

        x, ys = update
        x = np.asarray(x, dtype=np.float32).reshape([])
        ys = np.asarray(ys, dtype=np.float32).reshape([-1])
        ys_n = np.shape(ys)[0]

        if self._ys_n is None:
            self._ys_n = ys_n
        elif self._ys_n != ys_n:
            raise ValueError(f"Expected ys_n={self._ys_n}, got {ys_n}")

        self._xys_logs_buffer.extend(x.tobytes())
        self._xys_logs_buffer.extend(ys.tobytes())

    def _read_buffer(self) -> OutputType:
        if len(self._xys_logs_buffer) == 0 or self._ys_n is None:
            return None  # clear
        xys_logs = np.frombuffer(self._xys_logs_buffer, dtype=np.float32).reshape([-1, self._ys_n + 1])
        x_logs = xys_logs[:, 0]
        ys_logs = xys_logs[:, 1:]

        # histograms
        if self._bins is not None:
            histograms = [np.histogram(ys_logs[:, yi], bins=self._bins) for yi in range(self._ys_n)]
        else:
            histograms = None

        # curves
        self._cursor = bisect(x_logs, x_logs[-1] - self._win, lo=self._cursor)
        x_curve = x_logs[self._cursor:]
        ys_curve = ys_logs[self._cursor:]

        return x_curve, ys_curve, histograms


class LinesFigure(LiveFigure):
    def __init__(self, figure: Optional[plt.Figure] = None, *, win: float, hist: bool = False):
        super().__init__(figure)
        self._win = win
        self._hist = hist

        if self._hist:
            self.figure.set_size_inches(9, 4)
        else:
            self.figure.set_size_inches(8, 4)

        if self._hist:
            self._axes_curves, self._axes_histograms = self.figure.subplots(1, 2,
                sharey='row', gridspec_kw=dict(width_ratios=[3, 1], wspace=0.05))
            self._axes_curves.grid()
            self._axes_histograms.grid()
        else:
            self._axes_curves = self.figure.gca()
            self._axes_curves.grid()
            self._axes_histograms = None
        self._axes_curves: plt.Axes
        self._axes_histograms: Optional[plt.Axes]

        self._ys_lines: Optional[list[plt.Line2D]] = None
        self._ys_bars: Optional[list[Iterable[plt.Rectangle]]] = None

    @staticmethod
    def _init_lines(ys_n, axes_curves):
        return [axes_curves.plot(np.zeros([0]), zorder=2 + (ys_n - yi) / ys_n)[0]
                for yi in range(ys_n)]

    @staticmethod
    def _init_bars(ys_n, bars_n, axes_histograms):
        return [axes_histograms.bar(
            np.zeros([bars_n]), np.zeros([bars_n]),
            edgecolor="black", linewidth=1,
            zorder=2 + (ys_n - yi) / ys_n
        ) for yi in range(ys_n)]

    def _update_gui(self, data: OutputType):
        if data is None:  # clear
            self._axes_curves.cla()
            self._axes_curves.grid()
            self._ys_lines = None
            if self._axes_histograms:
                self._axes_histograms.cla()
                self._axes_histograms.grid()
                self._ys_bars = None
            return

        curves_x, curves_ys, ys_histograms = data
        if self._ys_lines is None:
            ys_n = np.shape(curves_ys)[-1]
            if ys_n == 0:
                return

            if self._hist:
                bars_n = len(ys_histograms[0][0])
                if bars_n == 0:
                    return
            else:
                bars_n = None

            self._ys_lines = self._init_lines(ys_n, self._axes_curves)
            self._ys_bars = self._init_bars(ys_n, bars_n, self._axes_histograms)

        # adjust y range
        min_y, max_y = np.min(curves_ys), np.max(curves_ys)
        if max_y > min_y:
            padding_y = (max_y - min_y) / 32
            self._axes_curves.set_ylim(min_y - padding_y, max_y + padding_y)
            if self._hist:
                self._axes_histograms.set_ylim(min_y - padding_y, max_y + padding_y)

        # update lines
        for yi, line in enumerate(self._ys_lines):
            line.set_xdata(curves_x)
            line.set_ydata(curves_ys[:, yi])

        # adjust x range
        min_x, max_x = np.min(curves_x), np.max(curves_x)
        if max_x < min_x + self._win:
            max_x = min_x + self._win
        else:
            min_x = max_x - self._win
        self._axes_curves.set_xlim(min_x, max_x)

        # update histogram
        if self._axes_histograms:
            max_height = 0
            for yi, (y_bars, (heights, edges)) in enumerate(zip(self._ys_bars, ys_histograms)):
                for yi, bar in enumerate(y_bars):
                    bar.set_y(edges[yi])
                    bar.set_height(edges[yi + 1] - edges[yi])
                    bar.set_width(heights[yi])
                max_height = max(max_height, np.max(heights))

            # adjust x range
            self._axes_histograms.set_xlim(0, max_height * 1.1)


class LiveLinesWidgets:
    def __init__(self,
            producer: Union[Iterator[tuple[float, Iterator[float]]], Callable[[], tuple[float, Iterator[float]]]], *,
            produce_interval: float = 1 / 128,
            frame_interval: float = 1 / 32,

            win: float = 200,
            bins: Optional[int] = None,

            figure: Optional[plt.Figure] = None,
            start: bool = True
    ):
        # ui widgets

        play_button = widgets.Button()
        play_button.description = "start"
        display(play_button)

        clear_button = widgets.Button()
        clear_button.description = "clear"
        display(clear_button)

        live_figure = LinesFigure(
            figure=figure,
            win=win,
            hist=bins is not None)

        # buffer

        live_buffer = LinesBufferChannel(
            interval=frame_interval,
            win=win,
            bins=bins)
        live_buffer.subscribe(live_figure)

        # producer

        if hasattr(producer, '__next__'):
            prod_func = lambda: next(producer)
        elif callable(producer):
            prod_func = producer
        else:
            raise TypeError(f"The producer is expected to be an iterator (generator) or callable, got {producer}")

        live_producer = IntervalLiveProducer(prod_func,
            interval=produce_interval)
        live_producer.subscribe(live_buffer)

        # use weak ref to avoid hanging by jupyter as displaying widgets

        producer_ref = weakref.ref(live_producer)
        buffer_ref = weakref.ref(live_buffer)

        # play button

        def on_play_button_clicked(_):
            producer = producer_ref()
            if producer is None:
                return
            play_button.disabled = True
            if producer.alive:
                producer.stop()
                play_button.description = "resume"
            else:
                producer.start()
                play_button.description = "pause"
            play_button.disabled = False

        play_button.on_click(on_play_button_clicked)

        # clear button

        def on_clear_button_clicked(_):
            buffer = buffer_ref()
            if buffer is None:
                return
            buffer.clear()

        clear_button.on_click(on_clear_button_clicked)

        # start
        if start:
            play_button.click()

        # properties
        self.play_button = play_button
        self.clear_button = clear_button
        self.producer = live_producer
        self.buffer = live_buffer
        self.figure = live_figure


def scan_lines_widgets(
        n, producer, *,
        sampling_rate=None,
        fps=64,
        win_size=200,
        hist_bins=None,
        figure=None,
        start=True
) -> tuple[IntervalLiveProducer, LinesFigure, widgets.Button]:
    warn(DeprecationWarning("Use cion.gui.live_lines instead."))
    live_lines_widgets = LiveLinesWidgets(
        producer,
        produce_interval=1 / sampling_rate,
        frame_interval=1 / fps,
        win=win_size,
        bins=hist_bins,
        figure=figure,
        start=start)
    return live_lines_widgets.producer, \
           live_lines_widgets.figure, \
           live_lines_widgets.play_button
